/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 *
 * File Name     : roce_tc.c
 * Version       : v2.0
 * Created       : 2024/5/15
 * Last Modified : 2024/5/15
 * Description   : The definition of RoCE DSCP functions.
 */

#include <rdma/ib_cache.h>
#include <linux/highmem.h>
#include <linux/inet.h>
#include <linux/sort.h>
#include <linux/xarray.h>


#include "roce.h"
#include "roce_qp.h"
#include "roce_mix.h"
#include "roce_cqm_cmd.h"
#include "roce_npu_cmd.h"
#include "hinic3_srv_nic.h"
#include "roce_tc.h"

#define MAX_MASK         32
#define GID_RAW_SHIFT    12
#define IPV4_ADDR_LENGTH 4
#define IPV6_ADDR_LENGTH 16

#if defined(OFED_MLNX_5_7_8) || defined(OFED_MLNX_5_4)
struct rdma_restrack_root {
    /**
     * @xa: Array of XArray structure to hold restrack entries.
     */
    struct xarray xa;
    /**
     * @next_id: Next ID to support cyclic allocation
     */
    u32 next_id;
};
#endif

static inline struct roce3_qp *to_mqp(struct ib_qp *ibqp)
{
    return container_of(ibqp, struct roce3_qp, ibqp);
}

/* roce3_force_tc feature */
static int check_string_match(const char *str, const char *str2)
{
    int str2_len;
    int str_len;

    if (!str || !str2) {
        return -EINVAL;
    }

    str_len = strlen(str);
    str2_len = strlen(str2);
    if (str_len <= str2_len) {
        return -EINVAL;
    }

    return memcmp(str, str2, str2_len);
}

static void tclass_set_mask_32(u32 *mask, int b)
{
    int bits = b;
    *mask = 0;
    if (!bits) {
        bits = MAX_MASK;
    }
    while (bits) {
        *mask = (*mask << 1) | 1;
        --bits;
    }
}

static int tclass_parse_src_ip(const char *str, void *store, void *store_mask)
{
    const char *end = NULL;

    return !in4_pton(str, -1, (u8 *)store, -1, &end);
}

static int tclass_parse_dst_ip(const char *str, void *store, void *store_mask)
{
    const char *end = NULL;
    int mask = 0;
    int ret;

    ret = !in4_pton(str, -1, (u8 *)store, -1, &end);
    if (ret) {
        pr_err("[ROCE] %s, in4_pton failed.ret: %d.\n", __func__, ret);
        return -EINVAL;
    }

    if (strlen(end)) {
        if (*end != '/') {
            pr_err("[ROCE] %s, parse dst ip failed.\n", __func__);
            return -EINVAL;
        }
        ret = kstrtoint(end + 1, 0, &mask);
        if (ret || mask < 0 || mask > MAX_MASK) {
            pr_err("[ROCE] %s, kstrtoint failed.ret: %d.\n", __func__, ret);
            return -EINVAL;
        }
    }

    tclass_set_mask_32(store_mask, mask);

    return ret;
}

static int tclass_parse_tclass(const char *str, void *ptr, void *store_mask)
{
    int *tclass = ptr;
    int ret;

    ret = kstrtoint(str, 0, tclass);
    if (ret || *tclass > 0xff) {
        return -EINVAL;
    }

    return 0;
}

static int tclass_compare_src_ips(struct tclass_match *match, struct tclass_match *match2, bool with_mask)
{
    return (int)(*(u32 *)match->s_addr != *(u32 *)match2->s_addr);
}

static int tclass_compare_dst_ips(struct tclass_match *match, struct tclass_match *match2, bool with_mask)
{
    u32 mask = 0xFFFFFFFF;

    if (with_mask) {
        mask = *(u32 *)match->d_addr_m;
    }

    return ((*(u32 *)match->d_addr & mask) != ((*(u32 *)match2->d_addr) & mask));
}

static size_t tclass_print_src_ip(struct tclass_match *match, char *buf, size_t size)
{
    return snprintf_s(buf, size, size - 1, "src_ip=%pI4,", match->s_addr);
}

static size_t tclass_print_dst_ip(struct tclass_match *match, char *buf, size_t size)
{
    return snprintf_s(buf, size, size - 1, "dst_ip=%pI4/%d,", match->d_addr,  hweight32(*(int *)match->d_addr_m));
}

static size_t tclass_print_tclass(struct tclass_match *match, char *buf, size_t size)
{
    return snprintf_s(buf, size, size - 1, "tclass=%d\n", match->tclass);
}

static const struct tclass_parse_node parse_tree[] = {
    TCLASS_CREATE_PARSE_NODE(TCLASS_MATCH_SRC_ADDR_IP, tclass_parse_src_ip, tclass_compare_src_ips,
        tclass_print_src_ip, "src_ip=", TCLASS_MATCH_MASK_SRC_ADDR_IP, s_addr, s_addr),
    TCLASS_CREATE_PARSE_NODE(TCLASS_MATCH_DST_ADDR_IP, tclass_parse_dst_ip, tclass_compare_dst_ips,
        tclass_print_dst_ip, "dst_ip=", TCLASS_MATCH_MASK_DST_ADDR_IP, d_addr, d_addr_m),
    TCLASS_CREATE_PARSE_NODE(TCLASS_MATCH_TCLASS, tclass_parse_tclass, NULL,
        tclass_print_tclass, "tclass=", TCLASS_MATCH_MASK_TCLASS, tclass, tclass),
    TCLASS_CREATE_PARSE_NODE(TCLASS_MATCH_TCLASS_NO_PREFIX, tclass_parse_tclass,
        NULL, NULL, "", TCLASS_MATCH_MASK_TCLASS, tclass, tclass),
};

static int tclass_verify_match(struct tclass_match *match)
{
    if (!(match->mask & TCLASS_MATCH_MASK_TCLASS)) {
        return -EINVAL;
    }

    if ((match->mask & (TCLASS_MATCH_MASK_SRC_ADDR_IP | TCLASS_MATCH_MASK_DST_ADDR_IP)) &&
        (match->mask & (TCLASS_MATCH_MASK_SRC_ADDR_IP6 | TCLASS_MATCH_MASK_DST_ADDR_IP6))) {
        return -EINVAL;
    }
    return 0;
}

static int tclass_parse_input_sub(struct tclass_match *match, char *p)
{
    int i;
    int ret;
    const struct tclass_parse_node *node = NULL;

    for (i = 0; i < ARRAY_SIZE(parse_tree); i++) {
        node = &parse_tree[i];
        if (!check_string_match(p, node->pattern)) {
            ret = parse_tree[i].parse(p + strlen(node->pattern),
                (char *)match + node->v_offset, (char *)match + node->m_offset);
            if (ret != 0) {
                return -EINVAL;
            }
            match->mask |= node->mask;
            break;
        }
    }
    if (i == ARRAY_SIZE(parse_tree)) {
        return -EINVAL;
    }
    return 0;
}

static int tclass_parse_input(char *str, struct tclass_match *match)
{
    char *p;
    int ret;

    while ((p = strsep(&str, ",")) != NULL) {
        if (!(*p)) {
            continue;
        }

        p = strim(p); /* Removing whitespace */
        ret = tclass_parse_input_sub(match, p);
        if (ret != 0) {
            return -EINVAL;
        }
    }

    return tclass_verify_match(match);
}

static struct tclass_match *tclass_find_empty(struct roce3_tc_data *tcd)
{
    int i;

    for (i = 0; i < TCLASS_MAX_RULES; i++) {
        if (!tcd->rule[i].mask) {
            return &tcd->rule[i];
        }
    }
        
    return NULL;
}

static struct tclass_match *tclass_find_match_sub(struct roce3_tc_data *tcd, struct tclass_match *match,
    u32 mask, bool with_mask, int i)
{
    int j;
    int ret = -1;
    for (j = 0; j < ARRAY_SIZE(parse_tree); j++) {
        const struct tclass_parse_node *node;

        node = &parse_tree[j];
        if ((mask & node->mask) && node->compare) {
            ret = node->compare(&tcd->rule[i], match, with_mask);
            if (ret != 0) {
                break;
            }
        }
    }
    if (!ret) {
        return &tcd->rule[i];
    }
    return NULL;
}

static struct tclass_match *tclass_find_match(struct roce3_tc_data *tcd, struct tclass_match *match,
    u32 mask, bool with_mask)
{
    int i;
    struct tclass_match *tag_match = NULL;

    mask |= TCLASS_MATCH_MASK_TCLASS;

    for (i = 0; i < TCLASS_MAX_RULES; i++) {
        if (tcd->rule[i].mask == mask) {
            tag_match = tclass_find_match_sub(tcd, match, mask, with_mask, i);
            if (tag_match != NULL) {
                return tag_match;
            }
        }
    }

    return NULL;
}

void tclass_get_tclass_locked(struct roce3_device *dev, struct roce3_tc_data *tcd, const struct rdma_ah_attr *ah,
    u8 port, u8 *tclass)
{
    struct tclass_match *res_match = NULL;
    struct tclass_match match = { 0 };
    enum ib_gid_type gid_type;
    union ib_gid gid;
    int mask;
    int err;

    if (tcd->val >= 0) {
        *tclass = tcd->val;
    } else if (ah && ah->type == RDMA_AH_ATTR_TYPE_ROCE) {
        err = rdma_query_gid(&dev->ib_dev, port, ah->grh.sgid_index, &gid);
        if (err) {
            goto out;
        }

        gid_type = ah->grh.sgid_attr->gid_type;
        if (gid_type != IB_GID_TYPE_ROCE_UDP_ENCAP) {
            goto out;
        }

        if (ipv6_addr_v4mapped((struct in6_addr *)&gid)) {
            match.mask = TCLASS_MATCH_MASK_SRC_ADDR_IP | TCLASS_MATCH_MASK_DST_ADDR_IP;
            (void)memcpy_s(match.s_addr, sizeof(match.s_addr), gid.raw + GID_RAW_SHIFT, IPV4_ADDR_LENGTH);
            (void)memcpy_s(match.d_addr, sizeof(match.d_addr), ah->grh.dgid.raw + GID_RAW_SHIFT, IPV4_ADDR_LENGTH);
        } else {
            match.mask = TCLASS_MATCH_MASK_SRC_ADDR_IP6 | TCLASS_MATCH_MASK_DST_ADDR_IP6;
            (void)memcpy_s(match.s_addr, IPV6_ADDR_LENGTH, gid.raw, IPV6_ADDR_LENGTH);
            (void)memcpy_s(match.d_addr, IPV6_ADDR_LENGTH, ah->grh.dgid.raw, IPV6_ADDR_LENGTH);
        }

        mask = match.mask;
        res_match = tclass_find_match(tcd, &match, mask, true);
        if (!res_match) {
            res_match = tclass_find_match(tcd, &match, mask &
                ~(TCLASS_MATCH_MASK_SRC_ADDR_IP | TCLASS_MATCH_MASK_SRC_ADDR_IP6), true);
        } else {
            goto out;
        }
        mask = match.mask;
        if (!res_match) {
            res_match = tclass_find_match(tcd, &match, mask &
                ~(TCLASS_MATCH_MASK_DST_ADDR_IP | TCLASS_MATCH_MASK_DST_ADDR_IP6), true);
        }
    }
out:
    if (res_match) {
        *tclass = res_match->tclass;
    }
}

void tclass_get_tclass(struct roce3_device *dev, struct roce3_tc_data *tcd, const struct rdma_ah_attr *ah,
    u8 port, u8 *tclass)
{
    mutex_lock(&tcd->lock);
    tclass_get_tclass_locked(dev, tcd, ah, port, tclass);
    mutex_unlock(&tcd->lock);
}
struct tc_attribute {
    struct attribute attr;
    ssize_t (*show)(struct roce3_tc_data *, struct tc_attribute *, char *buf);
    ssize_t (*store)(struct roce3_tc_data *, struct tc_attribute *, const char *buf, size_t count);
};

#define TC_ATTR(_name, _mode, _show, _store) \
    struct tc_attribute tc_attr_##_name = __ATTR(_name, _mode, _show, _store)

static ssize_t traffic_class_show(struct roce3_tc_data *tcd, struct tc_attribute *unused, char *buf)
{
    size_t count = 0;
    int j;
    int i;

    mutex_lock(&tcd->lock);
    if (tcd->val >= 0)
        count = snprintf_s(buf, PAGE_SIZE, PAGE_SIZE - 1, "Global tclass=%d\n", tcd->val);

    for (i = 0; i < TCLASS_MAX_RULES && count < (PAGE_SIZE - TCLASS_MAX_CMD); i++) {
        if (!tcd->rule[i].mask) {
            continue;
        }
        for (j = 0; j < ARRAY_SIZE(parse_tree); j++) {
            if ((tcd->rule[i].mask & parse_tree[j].mask) && parse_tree[j].print) {
                count += parse_tree[j].print(&tcd->rule[i], buf + count, PAGE_SIZE - count);
            }
        }
    }
    mutex_unlock(&tcd->lock);

    return count;
}

static int tclass_compare_match(const void *ptr1, const void *ptr2)
{
    const struct tclass_match *m1 = ptr1;
    const struct tclass_match *m2 = ptr2;

    if ((m1->mask & TCLASS_MATCH_MASK_DST_ADDR_IP) && (m2->mask & TCLASS_MATCH_MASK_DST_ADDR_IP)) {
        return hweight32(*(u32 *)m2->d_addr_m) - hweight32(*(u32 *)m1->d_addr_m);
    }

    if (m1->mask & TCLASS_MATCH_MASK_DST_ADDR_IP) {
        return -1;
    }

    if (m2->mask & TCLASS_MATCH_MASK_DST_ADDR_IP) {
        return 1;
    }

    return 0;
}

#if defined(OFED_MLNX_5_7_8) || defined(OFED_MLNX_5_4)
static int tclass_update_qp(struct roce3_device *ibdev, struct roce3_qp *mqp, u8 tclass)
{
    int ret;
    u8 cos;
    struct tag_cqm_cmd_buf *cqm_cmd_inbuf = NULL;
    roce_modify_dscp_value_s *modify_dscp_value = NULL;
    struct hinic3_dcb_state dcb = { 0 };

    ret = roce3_cqm_cmd_zalloc_inoutbuf(ibdev->hwdev, &cqm_cmd_inbuf,
        (u16)sizeof(roce_modify_dscp_value_s), NULL, 0);
    if (ret != 0) {
        dev_err(ibdev->hwdev_hdl, "[ROCE, ERR] %s: Failed to alloc cqm_cmd_inoutbuf, func_id(%u), ret(%d)\n",
            __func__, ibdev->glb_func_id, ret);
        return (-ENOMEM);
    }

    ret = hinic3_get_dcb_state(ibdev->hwdev, &dcb);
    if (ret != 0) {
        pr_err("[ROCE] %s, hinic3_get_dcb_state failed. ret: %d.\n", __func__, ret);
        roce3_cqm_cmd_free_inoutbuf(ibdev->hwdev, cqm_cmd_inbuf, NULL);
        return (-EINVAL);
    }

    modify_dscp_value = (roce_modify_dscp_value_s *)cqm_cmd_inbuf->buf;
    modify_dscp_value->com.index = cpu_to_be32(mqp->qpn);
    modify_dscp_value->com.dw0.bs.cmd_bitmask = cpu_to_be16(VERBS_CMD_TYPE_QP_BITMASK);
    modify_dscp_value->tcalss_value = tclass & 0xfc; // Bit 0 and bit 1 are ECN control bits. The tclass setting
                                                     // affects only the first six bits, that is, bit 7 to bit 2.
    if (dcb.trust == ROCE3_DCB_DSCP) {
        ret = hinic3_get_cos_by_pri(ibdev->hwdev, tclass >> ROCE3_DSCP_IDX, &cos);
        if (ret != 0) {
            pr_err("[ROCE] %s: get_cos_by_pri failed. ret: %d.\n", __func__, ret);
            roce3_cqm_cmd_free_inoutbuf(ibdev->hwdev, cqm_cmd_inbuf, NULL);
            return (-EINVAL);
        }
        modify_dscp_value->cos_value = htonl(cos);
    } else {
        modify_dscp_value->cos_value = -1;
    }

    ret = roce3_send_qp_lb_cmd(mqp->qpn, ibdev, ROCE_CMD_MODIFY_TCALSS_VALUE_QP,
        cqm_cmd_inbuf, NULL, ROCE_CMD_TIME_CLASS_A);
    if (ret != 0) {
        dev_err(ibdev->hwdev_hdl, "[ROCE, ERR] %s: Failed to send MODIFY_DSCP_VALUE_QP command, func_id(%u)\n",
            __func__, ibdev->glb_func_id);
        roce3_cqm_cmd_free_inoutbuf(ibdev->hwdev, cqm_cmd_inbuf, NULL);
        return -1;
    }
    pr_info("[ROCE] %s: Set tclass to qp. tclass = %u, cos = %u", __func__, tclass, cos);
    roce3_cqm_cmd_free_inoutbuf(ibdev->hwdev, cqm_cmd_inbuf, NULL);

    return 0;
}

static void set_tclass_to_qp(struct roce3_device *ibdev, struct roce3_tc_data *tcd, struct roce3_qp *mqp)
{
    int ret;
    u8 tclass;
    tclass = mqp->tclass;
    tclass_get_tclass_locked(ibdev, tcd, &mqp->ah, mqp->ah.port_num, &tclass);
    if (tclass != mqp->tclass) {
        ret = tclass_update_qp(ibdev, mqp, tclass);
        if (ret != 0) {
            return;
        }
        mqp->tclass = tclass;
    }
}
#endif

static void tclass_update_qps(struct roce3_tc_data *tcd)
{
#if defined(OFED_MLNX_5_7_8) || defined(OFED_MLNX_5_4)
    struct roce3_device *ibdev = tcd->ibdev;
    struct rdma_restrack_entry *res = NULL;
    struct rdma_restrack_root *rt = NULL;
    struct roce3_qp *mqp = NULL;
    unsigned long id = 0;
    struct ib_qp *ibqp = NULL;

    if (!tcd->ibdev) {
        return;
    }

    rt = &ibdev->ib_dev.res[RDMA_RESTRACK_QP];
    xa_lock(&rt->xa);
    xa_for_each(&rt->xa, id, res) {
        if (!rdma_restrack_get(res)) {
            continue;
        }

        xa_unlock(&rt->xa);

        ibqp = container_of(res, struct ib_qp, res);
        mqp = to_mqp(ibqp);
        if (ibqp->qp_type == IB_QPT_GSI || mqp->qp_type == IB_QPT_RESERVED4) {
            goto cont;
        }
        mutex_lock(&mqp->mutex);
        if (mqp->qp_state == IB_QPS_RTS && (rdma_ah_get_ah_flags(&mqp->ah) & IB_AH_GRH)) {
            set_tclass_to_qp(ibdev, tcd, mqp);
        }
        mutex_unlock(&mqp->mutex);
cont:
        rdma_restrack_put(res);
        xa_lock(&rt->xa);
    }
    xa_unlock(&rt->xa);
#endif
}

static ssize_t traffic_class_store(struct roce3_tc_data *tcd, struct tc_attribute *unused,
    const char *buf, size_t count)
{
    struct tclass_match *dst_match = NULL;
    char cmd[TCLASS_MAX_CMD + 1] = {};
    struct tclass_match match = {};
    int ret;

    if (count > TCLASS_MAX_CMD) {
        return -EINVAL;
    }
    ret = memcpy_s(cmd, sizeof(cmd), buf, count);
    if (ret != 0) {
        pr_err("[ROCE, ERR] %s: Memcpy failed. ret = %d", __func__, ret);
        return -EINVAL;
    }

    ret = tclass_parse_input(cmd, &match);
    if (ret) {
        return -EINVAL;
    }
    mutex_lock(&tcd->lock);
    if (match.mask == TCLASS_MATCH_MASK_TCLASS) {
        tcd->val = match.tclass;
    } else {
        dst_match = tclass_find_match(tcd, &match, match.mask, false);
        if (!dst_match) {
            dst_match = tclass_find_empty(tcd);
            if (!dst_match) {
                mutex_unlock(&tcd->lock);
                return -ENOMEM;
            }
        }
        if (match.tclass < 0) {
            memset_s(dst_match, sizeof(*dst_match), 0, sizeof(*dst_match));
        } else {
            memcpy_s(dst_match, sizeof(struct tclass_match), &match, sizeof(struct tclass_match));
        }
    }

    /* Sort the list based on subnet mask */
    sort(tcd->rule, TCLASS_MAX_RULES, sizeof(tcd->rule[0]), tclass_compare_match, NULL);
    tclass_update_qps(tcd);
    mutex_unlock(&tcd->lock);

    return count;
}

static TC_ATTR(traffic_class, 0644, traffic_class_show, traffic_class_store);

static struct attribute *tc_attrs[] = {
    &tc_attr_traffic_class.attr,
    NULL
};

static ssize_t tc_attr_show(struct kobject *kobj, struct attribute *attr, char *buf)
{
    struct tc_attribute *tc_attr = container_of(attr, struct tc_attribute, attr);
    struct roce3_tc_data *d = container_of(kobj, struct roce3_tc_data, kobj);

    if (!tc_attr->show) {
        return -EIO;
    }

    return tc_attr->show(d, tc_attr, buf);
}

static ssize_t tc_attr_store(struct kobject *kobj, struct attribute *attr, const char *buf, size_t count)
{
    struct tc_attribute *tc_attr = container_of(attr, struct tc_attribute, attr);
    struct roce3_tc_data *d = container_of(kobj, struct roce3_tc_data, kobj);

    if (!tc_attr->store) {
        return -EIO;
    }

    return tc_attr->store(d, tc_attr, buf, count);
}

static const struct sysfs_ops tc_sysfs_ops = {
    .show = tc_attr_show,
    .store = tc_attr_store
};

static struct kobj_type tc_type = {
    .sysfs_ops     = &tc_sysfs_ops,
    .default_attrs = tc_attrs
};

int init_tc_sysfs(struct roce3_device *dev)
{
    struct device *device = &dev->ib_dev.dev;
    int num_ports;
    int port;
    int err;

    dev->tc_kobj = kobject_create_and_add("tc", &device->kobj);
    if (!dev->tc_kobj) {
        return -ENOMEM;
    }
    num_ports = dev->hw_info.config_num_ports;
    for (port = 1; port <= num_ports; port++) {
        struct roce3_tc_data *tcd = &dev->tcd[port - 1];

        err = kobject_init_and_add(&tcd->kobj, &tc_type, dev->tc_kobj, "%d", port);
        if (err) {
            goto err;
        }
        tcd->val = -1;
        tcd->ibdev = dev;
        tcd->initialized = true;
        mutex_init(&tcd->lock);
    }
    return 0;
err:
    cleanup_tc_sysfs(dev);
    return err;
}

void cleanup_tc_sysfs(struct roce3_device *dev)
{
    if (dev->tc_kobj) {
        int num_ports;
        int port;

        kobject_put(dev->tc_kobj);
        dev->tc_kobj = NULL;
        num_ports = dev->hw_info.config_num_ports;
        for (port = 1; port <= num_ports; port++) {
            struct roce3_tc_data *tcd = &dev->tcd[port - 1];

            if (tcd->initialized) {
                kobject_put(&tcd->kobj);
            }
        }
    }
}